"""
麻将数据预处理器
"""
import os
import numpy as np
import pandas as pd
from typing import Dict, List, Optional, Tuple, Union
import json
import logging
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, MinMaxScaler

from sichuanmajiang.engine.cards import Tile, TileSet, Suit, convert_tile_string
from sichuanmajiang.engine.utils import (convert_hand_to_features,
                                          convert_tile_to_index,
                                          extract_game_features)


class DataPreprocessor:
    """
    麻将数据预处理器
    """
    def __init__(self,
                 feature_config: Optional[Dict] = None,
                 scaler_type: str = 'minmax'):  # 'minmax', 'standard', 'none'
        """
        Args:
            feature_config: 特征配置
            scaler_type: 缩放器类型
        """
        # 默认特征配置
        self.feature_config = feature_config or {
            'use_hand_features': True,
            'use_river_features': True,
            'use_tile_counts': True,
            'use_position_info': True,
            'use_score_info': True,
            'use_danger_scores': True,
            'use_expected_values': True
        }
        
        # 初始化缩放器
        self.scaler = None
        if scaler_type == 'minmax':
            self.scaler = MinMaxScaler()
        elif scaler_type == 'standard':
            self.scaler = StandardScaler()
        
        # 特征维度信息
        self.input_dim = self._calculate_input_dim()
        self.output_dim = 34  # 34种牌的决策
        
        # 初始化日志
        self.logger = logging.getLogger('DataPreprocessor')
    
    def _calculate_input_dim(self) -> int:
        """
        计算输入特征维度
        """
        dim = 0
        
        # 手牌特征 (34种牌 * 4张)
        if self.feature_config['use_hand_features']:
            dim += 34 * 4
        
        # 牌河特征 (34种牌)
        if self.feature_config['use_river_features']:
            dim += 34
        
        # 剩余牌数量 (34种牌)
        if self.feature_config['use_tile_counts']:
            dim += 34
        
        # 位置信息 (风向、轮次等)
        if self.feature_config['use_position_info']:
            dim += 8  # 简化的位置信息维度
        
        # 分数信息
        if self.feature_config['use_score_info']:
            dim += 4  # 4个玩家的分数差
        
        # 危险度信息
        if self.feature_config['use_danger_scores']:
            dim += 34
        
        # 期待值信息
        if self.feature_config['use_expected_values']:
            dim += 34
        
        return dim
    
    def preprocess_data(self, 
                       raw_data: List[Dict],
                       shuffle: bool = True,
                       validation_split: float = 0.2) -> Tuple[Dict, Dict, Dict]:
        """
        预处理原始数据
        
        Args:
            raw_data: 原始数据列表
            shuffle: 是否打乱数据
            validation_split: 验证集比例
            
        Returns:
            (训练数据, 验证数据, 测试数据) 字典
        """
        self.logger.info(f"开始预处理 {len(raw_data)} 条数据...")
        
        # 提取特征和标签
        features = []
        labels = []
        
        for data_point in raw_data:
            # 提取特征
            feature_vector = self._extract_features(data_point)
            features.append(feature_vector)
            
            # 提取标签
            label = self._extract_label(data_point)
            if label is not None:
                labels.append(label)
            else:
                # 如果没有标签，使用默认值
                labels.append(0)
        
        # 转换为numpy数组
        X = np.array(features)
        y = np.array(labels)
        
        # 数据清洗
        X, y = self._clean_data(X, y)
        
        # 划分数据集
        X_train, X_temp, y_train, y_temp = train_test_split(
            X, y, test_size=validation_split + 0.1, shuffle=shuffle, random_state=42
        )
        
        X_val, X_test, y_val, y_test = train_test_split(
            X_temp, y_temp, test_size=1/3, shuffle=shuffle, random_state=42
        )
        
        # 特征缩放
        if self.scaler:
            X_train = self.scaler.fit_transform(X_train)
            X_val = self.scaler.transform(X_val)
            X_test = self.scaler.transform(X_test)
        
        # 准备返回数据
        train_data = {
            'X': X_train,
            'y': y_train
        }
        
        val_data = {
            'X': X_val,
            'y': y_val
        }
        
        test_data = {
            'X': X_test,
            'y': y_test
        }
        
        self.logger.info(f"数据预处理完成！")
        self.logger.info(f"训练集: {len(X_train)} 样本, 验证集: {len(X_val)} 样本, 测试集: {len(X_test)} 样本")
        
        return train_data, val_data, test_data
    
    def _extract_features(self, data_point: Dict) -> np.ndarray:
        """
        从数据点提取特征向量
        
        Args:
            data_point: 数据点字典
            
        Returns:
            特征向量
        """
        features = []
        
        # 手牌特征
        if self.feature_config['use_hand_features'] and 'hand_features' in data_point:
            hand_features = np.array(data_point['hand_features'])
            features.append(hand_features.flatten())
        
        # 牌河特征
        if self.feature_config['use_river_features'] and 'river_features' in data_point:
            river_features = np.array(data_point['river_features'])
            features.append(river_features)
        
        # 剩余牌数量特征（简化版，基于牌河推断）
        if self.feature_config['use_tile_counts'] and 'river_features' in data_point:
            river_counts = np.array(data_point['river_features'])
            remaining_counts = 4 - river_counts  # 假设每种牌初始有4张
            features.append(remaining_counts)
        
        # 位置信息
        if self.feature_config['use_position_info']:
            position_features = self._extract_position_features(data_point)
            features.append(position_features)
        
        # 分数信息
        if self.feature_config['use_score_info'] and 'score' in data_point:
            score_features = self._extract_score_features(data_point)
            features.append(score_features)
        
        # 危险度信息
        if self.feature_config['use_danger_scores'] and 'danger_scores' in data_point:
            danger_scores = np.array(data_point['danger_scores'])
            features.append(danger_scores)
        
        # 期待值信息
        if self.feature_config['use_expected_values'] and 'expected_values' in data_point:
            expected_values = np.array(data_point['expected_values'])
            features.append(expected_values)
        
        # 合并所有特征
        if features:
            return np.concatenate(features)
        else:
            return np.zeros(self.input_dim)  # 返回默认特征向量
    
    def _extract_position_features(self, data_point: Dict) -> np.ndarray:
        """
        提取位置相关特征
        """
        # 简化的位置特征
        position_features = np.zeros(8)
        
        # 当前玩家的风向 (4个方向)
        if 'player_wind' in data_point:
            wind_map = {'EAST': 0, 'SOUTH': 1, 'WEST': 2, 'NORTH': 3}
            if data_point['player_wind'] in wind_map:
                position_features[wind_map[data_point['player_wind']]] = 1
        
        # 当前轮次的风向 (4个方向)
        if 'round_wind' in data_point:
            wind_map = {'EAST': 4, 'SOUTH': 5, 'WEST': 6, 'NORTH': 7}
            if data_point['round_wind'] in wind_map:
                position_features[wind_map[data_point['round_wind']]] = 1
        
        return position_features
    
    def _extract_score_features(self, data_point: Dict) -> np.ndarray:
        """
        提取分数相关特征
        """
        # 简化的分数特征
        score_features = np.zeros(4)
        
        # 当前玩家的分数
        if 'score' in data_point:
            # 标准化分数（假设分数范围在0-30000之间）
            normalized_score = min(data_point['score'] / 30000, 1.0)
            score_features[0] = normalized_score
        
        # 其他玩家的分数（这里使用占位符，实际应该从游戏状态获取）
        # 暂时设置为平均值
        score_features[1:] = 0.5
        
        return score_features
    
    def _extract_label(self, data_point: Dict) -> Optional[int]:
        """
        从数据点提取标签
        
        Args:
            data_point: 数据点字典
            
        Returns:
            标签（牌的索引）或None
        """
        if 'action' not in data_point or data_point['action'] is None:
            return None
        
        action = data_point['action']
        
        # 解析弃牌动作
        if action.startswith('discard_'):
            tile_str = action[8:]  # 去掉 'discard_' 前缀
            try:
                # 转换牌字符串为索引
                tile = convert_tile_string(tile_str)
                return convert_tile_to_index(tile)
            except:
                self.logger.warning(f"无法解析动作: {action}")
                return None
        
        # 对于其他动作（碰、杠、胡等），暂时返回-1
        return -1
    
    def _clean_data(self, X: np.ndarray, y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        清洗数据，移除无效样本
        
        Args:
            X: 特征数组
            y: 标签数组
            
        Returns:
            清洗后的特征和标签数组
        """
        # 移除标签为-1的样本
        valid_mask = y != -1
        X_clean = X[valid_mask]
        y_clean = y[valid_mask]
        
        # 移除包含NaN或无穷大的样本
        nan_mask = np.isnan(X_clean).any(axis=1)
        inf_mask = np.isinf(X_clean).any(axis=1)
        valid_mask = ~(nan_mask | inf_mask)
        
        X_clean = X_clean[valid_mask]
        y_clean = y_clean[valid_mask]
        
        # 记录清洗情况
        original_count = len(X)
        cleaned_count = len(X_clean)
        removed_count = original_count - cleaned_count
        
        if removed_count > 0:
            self.logger.info(f"数据清洗: 移除了 {removed_count} 个无效样本 ({removed_count/original_count*100:.2f}%)")
        
        return X_clean, y_clean
    
    def transform_state(self, game_state: Dict) -> np.ndarray:
        """
        转换游戏状态为模型输入特征
        
        Args:
            game_state: 游戏状态字典
            
        Returns:
            特征向量
        """
        # 提取特征
        feature_vector = self._extract_features(game_state)
        
        # 应用特征缩放
        if self.scaler:
            feature_vector = self.scaler.transform(feature_vector.reshape(1, -1))[0]
        
        return feature_vector
    
    def save_preprocessor(self, filepath: str) -> None:
        """
        保存预处理器配置
        
        Args:
            filepath: 保存路径
        """
        # 创建保存目录
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        
        # 保存配置
        config = {
            'feature_config': self.feature_config,
            'input_dim': self.input_dim,
            'output_dim': self.output_dim
        }
        
        # 保存配置到JSON
        config_file = filepath + '.json'
        with open(config_file, 'w', encoding='utf-8') as f:
            json.dump(config, f, ensure_ascii=False, indent=2)
        
        # 如果有缩放器，保存缩放器参数
        if self.scaler:
            if hasattr(self.scaler, 'scale_'):
                np.save(filepath + '_scale.npy', self.scaler.scale_)
            if hasattr(self.scaler, 'min_'):
                np.save(filepath + '_min.npy', self.scaler.min_)
            if hasattr(self.scaler, 'data_min_'):
                np.save(filepath + '_data_min.npy', self.scaler.data_min_)
            if hasattr(self.scaler, 'data_max_'):
                np.save(filepath + '_data_max.npy', self.scaler.data_max_)
            if hasattr(self.scaler, 'mean_'):
                np.save(filepath + '_mean.npy', self.scaler.mean_)
            if hasattr(self.scaler, 'var_'):
                np.save(filepath + '_var.npy', self.scaler.var_)
        
        self.logger.info(f"预处理器已保存到 {filepath}")
    
    def load_preprocessor(self, filepath: str) -> None:
        """
        加载预处理器配置
        
        Args:
            filepath: 加载路径
        """
        # 加载配置
        config_file = filepath + '.json'
        if not os.path.exists(config_file):
            self.logger.error(f"配置文件不存在: {config_file}")
            return
        
        with open(config_file, 'r', encoding='utf-8') as f:
            config = json.load(f)
        
        self.feature_config = config.get('feature_config', self.feature_config)
        self.input_dim = config.get('input_dim', self.input_dim)
        self.output_dim = config.get('output_dim', self.output_dim)
        
        # 加载缩放器参数
        if self.scaler:
            if hasattr(self.scaler, 'scale_') and os.path.exists(filepath + '_scale.npy'):
                self.scaler.scale_ = np.load(filepath + '_scale.npy')
            if hasattr(self.scaler, 'min_') and os.path.exists(filepath + '_min.npy'):
                self.scaler.min_ = np.load(filepath + '_min.npy')
            if hasattr(self.scaler, 'data_min_') and os.path.exists(filepath + '_data_min.npy'):
                self.scaler.data_min_ = np.load(filepath + '_data_min.npy')
            if hasattr(self.scaler, 'data_max_') and os.path.exists(filepath + '_data_max.npy'):
                self.scaler.data_max_ = np.load(filepath + '_data_max.npy')
            if hasattr(self.scaler, 'mean_') and os.path.exists(filepath + '_mean.npy'):
                self.scaler.mean_ = np.load(filepath + '_mean.npy')
            if hasattr(self.scaler, 'var_') and os.path.exists(filepath + '_var.npy'):
                self.scaler.var_ = np.load(filepath + '_var.npy')
        
        self.logger.info(f"预处理器已从 {filepath} 加载")


class FeatureEngineer:
    """
    麻将特征工程师
    """
    @staticmethod
    def create_hand_representation(hand: TileSet) -> np.ndarray:
        """
        创建手牌的表示形式
        
        Args:
            hand: 手牌集合
            
        Returns:
            手牌的二维数组表示 (27 x 4)
        """
        # 初始化27种牌，每种最多4张
        representation = np.zeros((27, 4), dtype=int)
        
        # 统计每种牌的数量
        tile_counts = {}
        for tile in hand:
            tile_idx = convert_tile_to_index(tile)
            if tile_idx not in tile_counts:
                tile_counts[tile_idx] = 0
            tile_counts[tile_idx] += 1
        
        # 填充表示矩阵
        for tile_idx, count in tile_counts.items():
            if count > 0:
                representation[tile_idx, :count] = 1
        
        return representation
    
    @staticmethod
    def create_3d_board_representation(game_state: Dict) -> np.ndarray:
        """
        创建游戏状态的3D表示形式（适合CNN模型）
        
        Args:
            game_state: 游戏状态字典
            
        Returns:
            3D表示 (channels, height, width)
        """
        # 创建5通道的表示：手牌、牌河、剩余牌、危险度、期待值
        channels = 5
        height = 4  # 4行：万、条、筒、字
        width = 9  # 9列
        
        representation = np.zeros((channels, height, width), dtype=float)
        
        # 1. 手牌通道
        if 'hand_tiles' in game_state:
            hand_tiles = [convert_tile_string(t) for t in game_state['hand_tiles']]
            for tile in hand_tiles:
                row, col = FeatureEngineer._get_tile_position(tile)
                if row < height and col < width:
                    representation[0, row, col] += 1/4  # 归一化
        
        # 2. 牌河通道
        if 'river_features' in game_state:
            river_features = np.array(game_state['river_features'])
            for idx in range(34):
                tile = FeatureEngineer._get_tile_from_index(idx)
                row, col = FeatureEngineer._get_tile_position(tile)
                if row < height and col < width:
                    representation[1, row, col] = river_features[idx] / 16  # 归一化
        
        # 3. 剩余牌通道
        # 简化计算：假设总共有136张牌
        total_tiles = 136
        used_tiles = sum(game_state.get('river_features', []) or [0]) * 4  # 简化计算
        remaining_ratio = (total_tiles - used_tiles) / total_tiles
        
        for idx in range(34):
            tile = FeatureEngineer._get_tile_from_index(idx)
            row, col = FeatureEngineer._get_tile_position(tile)
            if row < height and col < width:
                representation[2, row, col] = remaining_ratio
        
        # 4. 危险度通道
        if 'danger_scores' in game_state:
            danger_scores = np.array(game_state['danger_scores'])
            for idx in range(34):
                tile = FeatureEngineer._get_tile_from_index(idx)
                row, col = FeatureEngineer._get_tile_position(tile)
                if row < height and col < width:
                    representation[3, row, col] = danger_scores[idx]
        
        # 5. 期待值通道
        if 'expected_values' in game_state:
            expected_values = np.array(game_state['expected_values'])
            for idx in range(34):
                tile = FeatureEngineer._get_tile_from_index(idx)
                row, col = FeatureEngineer._get_tile_position(tile)
                if row < height and col < width:
                    representation[4, row, col] = expected_values[idx]
        
        return representation
    
    @staticmethod
    def _get_tile_position(tile: Tile) -> Tuple[int, int]:
        """
        获取牌在矩阵中的位置
        
        Args:
            tile: 牌对象
            
        Returns:
            (行, 列)
        """
        if tile.suit == Suit.CHARACTER:
            row = 0
            col = tile.rank - 1
        elif tile.suit == Suit.BAMBOO:
            row = 1
            col = tile.rank - 1
        elif tile.suit == Suit.DOT:
            row = 2
            col = tile.rank - 1
        else:  # 字牌
            row = 3
            # 东南西北白发中 -> 0-6
            if tile.suit == Suit.WIND:
                if tile.rank == 1:  # 东
                    col = 0
                elif tile.rank == 2:  # 南
                    col = 1
                elif tile.rank == 3:  # 西
                    col = 2
                else:  # 北
                    col = 3
            else:  # 三元牌
                if tile.rank == 1:  # 白
                    col = 4
                elif tile.rank == 2:  # 发
                    col = 5
                else:  # 中
                    col = 6
        
        return row, col
    
    @staticmethod
    def _get_tile_from_index(index: int) -> Tile:
        """
        从索引获取牌对象
        
        Args:
            index: 牌索引 (0-26)
            
        Returns:
            牌对象
        """
        # 简单实现，实际应该使用完整的转换逻辑
        # 这里仅作为示例
        if index < 9:
            return Tile(Suit.CHARACTER, index + 1)
        elif index < 18:
            return Tile(Suit.BAMBOO, index - 8)
        elif index < 27:
            return Tile(Suit.DOT, index - 17)
        raise ValueError(f"Invalid tile index: {index}")
    
    @staticmethod
    def calculate_combination_features(hand: TileSet) -> Dict[str, float]:
        """
        计算手牌组合特征
        
        Args:
            hand: 手牌集合
            
        Returns:
            组合特征字典
        """
        features = {}
        
        # 统计每种牌的数量
        tile_counts = {}
        for tile in hand:
            tile_idx = convert_tile_to_index(tile)
            tile_counts[tile_idx] = tile_counts.get(tile_idx, 0) + 1
        
        # 计算刻子数量
        triplet_count = sum(1 for count in tile_counts.values() if count >= 3)
        features['triplet_count'] = triplet_count
        
        # 计算对子数量
        pair_count = sum(1 for count in tile_counts.values() if count == 2)
        features['pair_count'] = pair_count
        
        # 计算单张数量
        singleton_count = sum(1 for count in tile_counts.values() if count == 1)
        features['singleton_count'] = singleton_count
        
        # 计算面子数量（刻子或顺子）
        # 这里使用简化计算
        features['面子数'] = triplet_count + (len(hand) - 3 * triplet_count - 2 * pair_count) / 3
        
        return features


def create_feature_pipeline(feature_config: Optional[Dict] = None) -> DataPreprocessor:
    """
    创建特征处理流水线
    
    Args:
        feature_config: 特征配置
        
    Returns:
        预处理器实例
    """
    return DataPreprocessor(feature_config=feature_config)